#include "svm_trainer.h"
#include <iostream>
#include <filesystem>
#include <random>

namespace fs = std::filesystem;

SVMTrainer::SVMTrainer() {
    // 构造函数
}

bool SVMTrainer::loadDataset(const std::string& dataset_path) {
    std::cout << "Loading dataset from: " << dataset_path << std::endl;
    
    // 清空现有数据
    all_features_ = cv::Mat();
    all_labels_ = cv::Mat();
    
    try {
        // 遍历 0-9 数字文件夹
        for (int label = 0; label <= 9; ++label) {
            std::string label_dir = dataset_path + "/train/" + std::to_string(label);
            
            if (!fs::exists(label_dir)) {
                std::cout << "Warning: Directory " << label_dir << " does not exist!" << std::endl;
                continue;
            }
            
            std::cout << "Loading images for number " << label << " from " << label_dir << std::endl;
            loadImagesFromDirectory(label_dir, label);
        }
        
        if (all_features_.rows == 0) {
            std::cout << "No training data found!" << std::endl;
            return false;
        }
        
        std::cout << "Total samples loaded: " << all_features_.rows << std::endl;
        std::cout << "Feature dimensions: " << all_features_.cols << std::endl;
        
        return true;
        
    } catch (const fs::filesystem_error& ex) {
        std::cout << "Filesystem error: " << ex.what() << std::endl;
        return false;
    }
}

void SVMTrainer::loadImagesFromDirectory(const std::string& dir_path, int label) {
    int loaded_count = 0;
    
    for (const auto& entry : fs::directory_iterator(dir_path)) {
        if (entry.is_regular_file()) {
            std::string extension = entry.path().extension().string();
            if (extension == ".jpg" || extension == ".png" || extension == ".bmp") {
                cv::Mat image = cv::imread(entry.path().string(), cv::IMREAD_GRAYSCALE);
                if (!image.empty()) {
                    cv::Mat features = preprocessImage(image);
                    
                    if (all_features_.empty()) {
                        all_features_ = features.clone();
                    } else {
                        cv::vconcat(all_features_, features, all_features_);
                    }
                    
                    cv::Mat label_mat = (cv::Mat_<float>(1, 1) << static_cast<float>(label));
                    if (all_labels_.empty()) {
                        all_labels_ = label_mat.clone();
                    } else {
                        cv::vconcat(all_labels_, label_mat, all_labels_);
                    }
                    
                    loaded_count++;
                }
            }
        }
    }
    
    std::cout << "Loaded " << loaded_count << " images for label " << label << std::endl;
}

cv::Mat SVMTrainer::preprocessImage(const cv::Mat& image) {
    cv::Mat processed;
    
    // 调整大小到 28x28
    cv::resize(image, processed, cv::Size(28, 28));
    
    // 二值化
    cv::threshold(processed, processed, 0, 255, cv::THRESH_BINARY | cv::THRESH_OTSU);
    
    // 转换为浮点数并归一化
    processed.convertTo(processed, CV_32F, 1.0 / 255.0);
    
    // 展平为一行
    return processed.reshape(1, 1);
}

void SVMTrainer::setTrainTestSplit(double train_ratio) {
    if (all_features_.rows == 0) {
        std::cout << "No data to split!" << std::endl;
        return;
    }
    
    int total_samples = all_features_.rows;
    int train_samples = static_cast<int>(total_samples * train_ratio);
    
    // 创建随机索引
    std::vector<int> indices(total_samples);
    for (int i = 0; i < total_samples; ++i) {
        indices[i] = i;
    }
    
    // 随机打乱
    std::random_device rd;
    std::mt19937 g(rd());
    std::shuffle(indices.begin(), indices.end(), g);
    
    // 分割数据
    train_features_ = cv::Mat(train_samples, all_features_.cols, CV_32F);
    train_labels_ = cv::Mat(train_samples, 1, CV_32F);
    test_features_ = cv::Mat(total_samples - train_samples, all_features_.cols, CV_32F);
    test_labels_ = cv::Mat(total_samples - train_samples, 1, CV_32F);
    
    for (int i = 0; i < train_samples; ++i) {
        all_features_.row(indices[i]).copyTo(train_features_.row(i));
        all_labels_.row(indices[i]).copyTo(train_labels_.row(i));
    }
    
    for (int i = train_samples; i < total_samples; ++i) {
        all_features_.row(indices[i]).copyTo(test_features_.row(i - train_samples));
        all_labels_.row(indices[i]).copyTo(test_labels_.row(i - train_samples));
    }
    
    std::cout << "Data split: " << train_samples << " training, " 
              << (total_samples - train_samples) << " testing" << std::endl;
}

bool SVMTrainer::trainModel(const std::string& model_save_path) {
    if (train_features_.rows == 0) {
        std::cout << "No training data available!" << std::endl;
        return false;
    }
    
    std::cout << "Training SVM model..." << std::endl;
    std::cout << "Training set: " << train_features_.rows << " samples" << std::endl;
    
    // 创建SVM模型
    cv::Ptr<cv::ml::SVM> svm = cv::ml::SVM::create();
    
    // 设置SVM参数
    svm->setType(cv::ml::SVM::C_SVC);
    svm->setKernel(cv::ml::SVM::LINEAR);
    svm->setTermCriteria(cv::TermCriteria(cv::TermCriteria::MAX_ITER + cv::TermCriteria::EPS, 1000, 1e-6));
    
    // 设置C参数（正则化参数）
    svm->setC(1.0);
    
    // 训练模型
    try {
        std::cout << "Starting training..." << std::endl;
        svm->train(train_features_, cv::ml::ROW_SAMPLE, train_labels_);
        std::cout << "Training completed!" << std::endl;
        
        // 评估模型
        double train_accuracy = evaluateModel(svm);
        std::cout << "Training accuracy: " << train_accuracy * 100 << "%" << std::endl;
        
        // 保存模型
        svm->save(model_save_path);
        std::cout << "Model saved to: " << model_save_path << std::endl;
        
        return true;
        
    } catch (const cv::Exception& e) {
        std::cout << "Error during training: " << e.what() << std::endl;
        return false;
    }
}

double SVMTrainer::evaluateModel(const cv::Ptr<cv::ml::SVM>& svm) {
    if (test_features_.rows == 0) {
        std::cout << "No test data available for evaluation!" << std::endl;
        return 0.0;
    }
    
    int correct_predictions = 0;
    int total_predictions = test_features_.rows;
    
    for (int i = 0; i < total_predictions; ++i) {
        cv::Mat sample = test_features_.row(i);
        float true_label = test_labels_.at<float>(i, 0);
        float predicted_label = svm->predict(sample);
        
        if (std::abs(predicted_label - true_label) < 0.1) {
            correct_predictions++;
        }
    }
    
    double accuracy = static_cast<double>(correct_predictions) / total_predictions;
    std::cout << "Test accuracy: " << accuracy * 100 << "% (" 
              << correct_predictions << "/" << total_predictions << ")" << std::endl;
    
    return accuracy;
}